'''
@Project ：dataProcessing.py 
@File    ：SVMModel.py
@IDE     ：PyCharm 
@Author  ：子协
@Date    ：2025-03-28 18:27 
'''
from libsvm.svmutil import *
from libsvm.svm import *
from libsvm.commonutil import *


class SVMModel:
    def __init__(self, kernel="rbf", C=2048.0, gamma=0.0):
        kernel_dict = {"linear": 0, "poly": 1, "rbf": 2, "sigmoid": 3, "precomputed": 4}
        self.model = None
        self.kernel = kernel_dict[kernel]
        self.C = C
        self.gamma = gamma

    def fit(self, x, y):
        prob = svm_problem(y, x)
        param = svm_parameter("-t {} -c {} -g {}".format(self.kernel,self.C,self.gamma))
        self.model = svm_train(prob, param)

    def predict(self, *args):
        if self.model is None:
            raise ValueError("Model not trained yet.")
        if len(args) == 1:
            return svm_predict([], args[0], self.model)
        elif len(args) == 2:
            return svm_predict(args[1], args[0], self.model)
        else:
            raise ValueError("The maximum number of parameters is 2.")

    def saveModel(self,path):
        svm_save_model(path, self.model)

    def loadModel(self,path):
        self.model = svm_load_model(path)

